

#include <stdlib.h>
#include <stdio.h>
#include <string.h>
#include <math.h>

#include <iostream>

#include "mrilib.h"

using std::cout;
using std::cerr;
using std::endl;

int main (int argc, char **argv)
{
	int *imsize = 0, nimsize = 0; 
	int pm_iters = 2;
	int qcols = 0;
	shm_array_t *shm_KKt = 0;

	struct arg* args [] = {
		new arg_t<int> ("pm_iters",  pm_iters),
		new arg_t<int> ("qcols",  qcols),
		new list_arg_t<int> ("imsize", imsize, nimsize, "x"),
		0
	};

	get_args (argc, argv, args);

	if (!imsize || nimsize != 3){
		cerr << "Must specify imsize=<int>x<int>x<int>" << endl;
		return -1;
	}

	int px = imsize[0], py = imsize[1], ch = imsize[2];

	if (qcols <= 0 || qcols > ch)
		qcols = ch;

	if (get_shm_array ("KKt", &shm_KKt)){
		perror("KKt");
		return -1;
	}

	int ks1 = shm_KKt->dims[0], ks2 = shm_KKt->dims[1];

	if (shm_KKt->ndims != 4 or (shm_KKt->dims[2] != shm_KKt->dims[3])
	      or (shm_KKt->dims[2] != ch)){
		cerr << "KKt must be <ks1>x<ks2>x<ch>x<ch>!" << endl;
		return -1;
	}

	if (shm_KKt->type != shm_array_cplxf){
		cerr << "KKt must be single-precision complex!" << endl;
		return -1;
	}
	complex *KKt = (complex*) shm_KKt->data.c;

	complex *km = new complex [px*py*ch*ch];
	for (int i = 0; i < px*py*ch*ch; ++i){ km[i].real = km[i].imag = 0.0f; }

	int o1 = pad_index (px, ks1), o2 = pad_index (py, ks2);

	for (int c2 = 0; c2 < ch; ++c2)
		for (int c1 = 0; c1 < ch; ++c1)
			for (int k2 = 0; k2 < ks2; ++k2)
				for (int k1 = 0; k1 < ks1; ++k1)
				{
					complex &KKt_ = KKt [ks1-1-k1 + ks1*(ks2-1-k2 + ks2*(c1 + ch*c2))];
					complex &km_ = km [o1+k1 + px*(o2+k2 + py*(c1 + ch*c2))];

					km_ = complex_multiply (compose_complex (sqrt(px*py), 0), KKt_);
				}

	fftwf_plan forward, backward;
	plan_fft2c (forward, backward, (cplxf*) km, px, py, ch*ch);
	fft2c_exec (backward, 1.0 / (px*py), (cplxf*) km, px, py, ch*ch);
	destroy_fft2c (forward, backward);

	complex *smap = new complex [px*py*ch*qcols];

	{
		acmlsetnumthreads (1);

		complex *A = new complex [ch*ch],
		        *Q = new complex [ch*qcols],
		        *T = new complex [ch*qcols],
		        *v = new complex [qcols];

		complex _1 [1] = {{1.0f, 0.0f}}, _0 [1] = {{0.0f, 0.0f}},
		        _n1[1] = {{-1.0f, 0.0f}};

		for (int iy = 0; iy < py; ++iy)
			for (int ix = 0; ix < px; ++ix)
			{
				for (int c2 = 0; c2 < ch; ++c2)
					for (int c1 = 0; c1 < ch; ++c1)
						A [c1 + ch*c2] = km [ix + px*(iy + py*(c1 + ch*c2))];

				for (int c2 = 0; c2 < qcols; ++c2)
					for (int c1 = 0; c1 < ch; ++c1)
						Q [c1 + ch*c2].real = Q [c1 + ch*c2].imag = 0;

				for (int c1 = 0; c1 < qcols; ++c1)
					Q [c1 + ch*c1].real = 1.0f;

				for (int iter = 1; iter <= pm_iters; ++iter)
				{
					cgemm ('N', 'N', ch, qcols, ch, _1, A, ch, Q, ch, _0, T, ch);

					for (int j = 0; j < qcols; ++j){
						complex nrm = {0,0};
						nrm.real = 1.0f / scnrm2 (ch, &T [ch*j], 1);
						cscal (ch, &nrm, &T [ch*j], 1);

					/*
						float s = scnrm2 (ch, &T [ch*j], 1);
						for (int i = 0; i < ch; ++i){
							T[i + ch*j].real /= s;
							T[i + ch*j].imag /= s;
						}
					*/
			
					}

					for (int j = 1; j < qcols; ++j)
					{
						cgemv  ('C', ch,j, _n1, T,ch, &T[j*ch],1, _0, v,1);

						for (int k = 0; k < j; ++k)
							caxpy (ch, &v[k], &T[k*ch],1, &T[j*ch],1);
					}

					complex *tmp = Q; Q = T; T = tmp;
				}

				for (int c2 = 0; c2 < qcols; ++c2)
					for (int c1 = 0; c1 < ch; ++c1)
						smap [ix + px*(iy + py*(c1 + ch*c2))] = Q [c1 + ch*c2];
			}
	}

	shm_array_t sz_smap, *shm_smap;

	sz_smap.type = shm_array_cplxf;
	sz_smap.ndims = 4;
	sz_smap.dims[0] = px;
	sz_smap.dims[1] = py;
	sz_smap.dims[2] = ch;
	sz_smap.dims[3] = qcols;

	make_shm_array ("smap", &sz_smap, &shm_smap);
	memcpy ((void*) shm_smap->data.c, (void*) smap, px*py*ch*qcols *sizeof(*smap));
	put_shm_array ("smap", shm_smap);

	return 0;
}
